import tensorflow as tf
import re
import os
from tensorflow.keras import layers, losses
from eroica.keras import get_train_val, get_all
from eroica.keras import TextVectorization  # 引用文本预处理层

data_url = r"D:\casket\Keras_data\stack_overflow_16k"
train_url = os.path.join(data_url, 'train')  # 获取训练集路径
test_url = os.path.join(data_url, 'test')  # 获取测试集路径

# 获取训练集和验证集数据
train_data, val_data = get_train_val(train_url)
# 获取测试集数据
test_data = get_all(test_url)

# 预处理数据，把文本数据变成数字序列
list_len = 10000  # 词表大小
str_len = 250  # 序列长度

# 文本预处理层
vectorize_layer = TextVectorization(list_len, str_len)

# 训练词表
train_text = train_data.map(lambda x, y: x)
vectorize_layer.adapt(train_text)


# 定义把字符串转换为数字序列的函数
def vectorize_text(text, label):
    text = tf.expand_dims(text, -1)
    return vectorize_layer(text), label


train_ds = train_data.map(vectorize_text)
val_ds = val_data.map(vectorize_text)
test_ds = test_data.map(vectorize_text)

# 性能优化
AUTOTUNE = tf.data.AUTOTUNE  # 框架会根据硬件自动调整要取多少个bath合适。
train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)
test_ds = test_ds.cache().prefetch(buffer_size=AUTOTUNE)

# 构建模型
model = tf.keras.Sequential([
    lears.Embedding(list_len + 1, 64, input_length=str_len),
    layers.GlobalAveragePooling1D(),
    layers.Dense(4),
])
